from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives import serialization, hashes


def encrypt_with_public_key(public_key_pem, data):
    """
    使用公钥对数据进行加密

    :param public_key_pem: 公钥的 PEM 格式字符串
    :param data: 要加密的数据，字符串类型
    :return: 加密后的字节数据
    """
    # 从 PEM 格式的公钥字符串中加载公钥对象
    public_key = serialization.load_pem_public_key(
        public_key_pem.encode('utf-8'),
        backend=default_backend()
    )

    # 将待加密的数据编码为字节类型
    # data_bytes = data.encode('utf-8')

    # 使用公钥对数据进行加密，采用 OAEP 填充和 SHA256 哈希算法
    ciphertext = public_key.encrypt(
        # data_bytes,
        padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None
        )
    )

    return ciphertext


if __name__ == "__main__":
    # 示例公钥，实际使用时需替换为真实的公钥
    public_key_pem = """
-----BEGIN PUBLIC KEY-----
MIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQCc4jyuRP/Z7ObB7dUp9joHE61fnBXsMzC/frdNiTmxilSYuEx/owDdijPGIHP7xJqwHk+5tI3n9SDgyVb9eeZhymaCpSGt2fZoz2O6XrNKme3TL29WH9/0oMzqQnCFvNWNP+FYnbY9uOYafhRhWWRNP8+kb2mXCSsplBNHTTa77QIDAQAB
-----END PUBLIC KEY-----
"""

    # 要加密的数据
    data_to_encrypt = {"f56": "华研检测集团有限责任公司", "f1": "2024090508", "f79": ""}

    # 调用加密函数进行加密
    encrypted_data = encrypt_with_public_key(public_key_pem, data_to_encrypt)

    print("加密后的数据（字节形式）:", encrypted_data)
